import os
import sys
import gradio as gr
import numpy as np
import torch
import cv2
from PIL import Image
import time

# 设置项目路径
current_dir = os.path.dirname(os.path.abspath(__file__))
segment_anything_path = os.path.join(current_dir, "segment-anything-2")
sys.path.append(segment_anything_path)

from sam2.build_sam import build_sam2
from sam2.sam2_image_predictor import SAM2ImagePredictor

# 初始化模型
checkpoint_path = os.path.join(segment_anything_path, "checkpoints", "sam2.1_hiera_tiny.pt")
config_path = os.path.join(segment_anything_path, "sam2", "configs", "sam2.1", "sam2.1_hiera_t.yaml")

try:
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = build_sam2(config_path, checkpoint_path)
    model.to(device)
    predictor = SAM2ImagePredictor(model)
    print(f"模型已加载到设备: {device}")
except Exception as e:
    print(f"模型加载失败: {str(e)}")

    sys.exit(1)

# 全局变量
points_list = []
original_image = None

def on_image_change(image):
    """
    当用户上传或更改图片时，初始化全局变量
    """
    global points_list, original_image
    points_list = []
    if image is not None:
        original_image = image.copy()
    return

def on_click(evt: gr.SelectData):
    """
    处理图片点击事件，记录点击的位置并更新显示
    """
    global points_list, original_image
    if evt is not None and original_image is not None:
        x, y = evt.index
        points_list.append([x, y])
        print(f"点击位置: ({x}, {y})")

        # 绘制点击位置在图像上
        temp_image = original_image.copy()
        for point in points_list:
            cv2.circle(temp_image, (point[0], point[1]), radius=5, color=(0, 255, 0), thickness=-1)
        return temp_image
    return original_image

def clear_points():
    """

    清除所有已记录的点并重置图像
    """
    global points_list, original_image
    points_list = []
    print("已清除所有点")
    if original_image is not None:
        return original_image
    else:
        return None

def segment_image(image):
    """
    处理图像分割并保存结果
    """
    global points_list, predictor

    if image is None:
        print("错误：未上传图片")
        return None

    if len(points_list) == 0:
        print("错误：未选择点击位置")
        return None

    try:
        # 确保图像格式正确
        if isinstance(image, np.ndarray):
            if len(image.shape) == 2:  # 灰度图转RGB
                image = cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
            elif image.shape[2] == 4:  # RGBA转RGB
                image = image[..., :3]

        # 设置图像
        predictor.set_image(image)

        # 预测分割结果
        input_points = np.array(points_list)
        input_labels = np.ones(len(points_list))

        masks, _, _ = predictor.predict(
            point_coords=input_points,
            point_labels=input_labels,
            multimask_output=False
        )

        # 处理预测结果
        mask = masks[0].astype(np.uint8) * 255

        # 创建带透明通道的输出图像
        output = np.zeros((image.shape[0], image.shape[1], 4), dtype=np.uint8)
        output[..., :3] = image
        output[..., 3] = mask

        # 转换为PIL图像
        output_image = Image.fromarray(output, 'RGBA')

        # 保存结果
        timestamp = time.strftime("%Y%m%d_%H%M%S")
        save_dir = "output"
        os.makedirs(save_dir, exist_ok=True)
        save_path = os.path.join(save_dir, f"segment_result_{timestamp}.png")
        output_image.save(save_path)

        print(f"结果已保存至: {save_path}")

        # 清除点击记录，为下一次分割做准备
        points_list.clear()

        return output_image

    except Exception as e:
        print(f"处理过程中出错: {str(e)}")
        return None

# 创建Gradio界面
with gr.Blocks() as demo:
    gr.Markdown("# SAM2图像分割工具")

    with gr.Row():
        with gr.Column():
            input_image = gr.Image(
                label="上传图片并点击目标位置",
                type="numpy",
                interactive=True
            )
            clear_btn = gr.Button("清除点击", variant="secondary")
            segment_btn = gr.Button("开始分割", variant="primary")
        with gr.Column():
            click_image = gr.Image(
                label="标记后的图像",
                type="numpy"
            )
            output_image = gr.Image(
                label="分割结果",
                type="pil"
            )

    # 设置事件触发
    input_image.change(
        fn=on_image_change,
        inputs=[input_image],
        outputs=None
    )

    input_image.select(
        fn=on_click,
        inputs=None,
        outputs=[click_image]
    )

    clear_btn.click(
        fn=clear_points,
        inputs=None,
        outputs=[click_image]
    )

    segment_btn.click(
        fn=segment_image,
        inputs=[input_image],
        outputs=[output_image]
    )

# 启动应用
if __name__ == "__main__":
    demo.launch(
        server_name="127.0.0.1",
        server_port=7860,
        share=True
    )
